/*
 * To change this template, choose Tools | Templates
 * and open the template in the editor.
 */
package server.votesending.shamirsecret;

import java.io.BufferedInputStream;
import java.io.BufferedOutputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
import java.math.BigInteger;
import java.security.InvalidKeyException;
import java.security.Key;
import java.security.KeyFactory;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.NoSuchAlgorithmException;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.interfaces.RSAPrivateKey;
import java.security.interfaces.RSAPublicKey;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.RSAPrivateKeySpec;
import java.security.spec.RSAPublicKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.Random;
import java.util.logging.Level;
import java.util.logging.Logger;
import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.NoSuchPaddingException;

/**
 *
 * @author Maciek
 */
public class GenerateKeysAndSecret {

    RSAPublicKey publicKey;
    RSAPrivateKey privateKey;
    KeyPair kp;

    public RSAPrivateKey getPrivateKey() {
        return privateKey;
    }

    public void setPrivateKey(RSAPrivateKey privateKey) {
        this.privateKey = privateKey;
    }

    public RSAPublicKey getPublicKey() {
        return publicKey;
    }

    public void setPublicKey(RSAPublicKey publicKey) {
        this.publicKey = publicKey;
    }

    public void generateKeys(int keylength) throws NoSuchAlgorithmException {
        KeyPairGenerator kpg = KeyPairGenerator.getInstance("RSA");
        kpg.initialize(2048);
        kp = kpg.genKeyPair();
        publicKey = (RSAPublicKey) kp.getPublic();
        privateKey = (RSAPrivateKey) kp.getPrivate();


    }

    public void saveKeys(String path) throws NoSuchAlgorithmException, InvalidKeySpecException, IOException {
        storeKeys(path, kp);
    }

    public String generateRandomPassword(int len) {
        StringBuilder sb = new StringBuilder();
        Random random = new Random();
        for (int i = 0; i < len; ++i) {
            sb.append(random.nextInt(10));
        }
        return sb.toString();

    }

    public void storeKeysX509(String path, int passLen, int k, int n, String base) throws FileNotFoundException, IOException, BadPaddingException, InterruptedException {

        writeBytesToFile(publicKey.getEncoded(), path + File.separatorChar + VotingKeyPair.publicKeyFile);
        System.out.println("Dlugosc pubkey : " + publicKey.getEncoded().length);

        CypherModel cyphermodel = new CypherModel();
        String password = generateRandomPassword(passLen);
        System.out.println("passLen : " + passLen + " password : " + password);
        PKCS8EncodedKeySpec pkcs8EncodedKeySpec = new PKCS8EncodedKeySpec(
                privateKey.getEncoded());
        byte[] encryptedData = cyphermodel.encrypt(pkcs8EncodedKeySpec.getEncoded(), password);
        if (encryptedData == null) {
            System.out.println("Nie udało się zaszyfrować klucza prywatnego ");
            return;
        }
        writeBytesToFile(encryptedData, path + File.separatorChar + VotingKeyPair.privateKeyFile);

        ShamirsSecret ss = new ShamirsSecret();
        ss.fileShamirSecret(path + File.separatorChar + base, password, k, n);
    }

    public void restoreKeysX509(String path, int k, int ids[], String base) throws FileNotFoundException, IOException, NoSuchAlgorithmException, InvalidKeySpecException, InterruptedException, BadPaddingException {
        try {
            byte[] pubKey = readBytesFromFile(path + File.separatorChar + VotingKeyPair.publicKeyFile);
            // byte[] pubKey = publicKey.getEncoded();

            X509EncodedKeySpec pubKeySpec = new X509EncodedKeySpec(pubKey);
            KeyFactory keyFactory = KeyFactory.getInstance("RSA");
            RSAPublicKey nPublicKey = (RSAPublicKey) keyFactory.generatePublic(pubKeySpec);


            // public key zrobiona 
            CypherModel cyphermodel = new CypherModel();
            ShamirsSecret ss = new ShamirsSecret();
            int result = ss.fileCombineShamirSecret(base, k, ids);
            if (result != 0) {
                System.out.println("Tajemnica shamira nie wyszła");
            }
            String password = ss.getSecret();
            System.out.println("haslo  " + password);
            byte[] privKey = cyphermodel.decryptToBytes(readBytesFromFile(path + File.separatorChar + VotingKeyPair.privateKeyFile), password);

            PKCS8EncodedKeySpec privateKeySpec = new PKCS8EncodedKeySpec(
                    privKey);

            RSAPrivateKey nPrivateKey = (RSAPrivateKey) keyFactory.generatePrivate(privateKeySpec);

            /*
             *  Tu będzie test czy się zgadza:
            
            
            System.out.println("Write public key : " + nPublicKey.getPublicExponent() + "\n" + nPublicKey.getModulus());
            System.out.println("Write private key : " + nPrivateKey.getPrivateExponent() + "\n" + nPrivateKey.getModulus());
            
            System.out.println("New publickey okey : " + nPublicKey.equals(publicKey));
            System.out.println("New privatekey okey : " + nPrivateKey.equals(privateKey));
            
            System.out.println("Public exponents equal : " + nPublicKey.getPublicExponent().equals(publicKey.getPublicExponent()));
            System.out.println("Private exponents equal : " + nPrivateKey.getPrivateExponent().equals(privateKey.getPrivateExponent()));
             * 
             */

            privateKey = nPrivateKey;
            publicKey = nPublicKey;

            this.checkCommunication();
        } catch (NoSuchPaddingException ex) {
            Logger.getLogger(GenerateKeysAndSecret.class.getName()).log(Level.SEVERE, null, ex);
        } catch (InvalidKeyException ex) {
            Logger.getLogger(GenerateKeysAndSecret.class.getName()).log(Level.SEVERE, null, ex);
        } catch (IllegalBlockSizeException ex) {
            Logger.getLogger(GenerateKeysAndSecret.class.getName()).log(Level.SEVERE, null, ex);
        }
    }

    public void restorePubKeyX509(String path) {
        try {
            byte[] pubKey = readBytesFromFile(path + File.separatorChar + VotingKeyPair.publicKeyFile);
            // byte[] pubKey = publicKey.getEncoded();

            X509EncodedKeySpec pubKeySpec = new X509EncodedKeySpec(pubKey);
            KeyFactory keyFactory = KeyFactory.getInstance("RSA");
            RSAPublicKey nPublicKey = (RSAPublicKey) keyFactory.generatePublic(pubKeySpec);
            publicKey = nPublicKey;
        } catch (InvalidKeySpecException ex) {
            Logger.getLogger(GenerateKeysAndSecret.class.getName()).log(Level.SEVERE, null, ex);
        } catch (NoSuchAlgorithmException ex) {
            Logger.getLogger(GenerateKeysAndSecret.class.getName()).log(Level.SEVERE, null, ex);
        } catch (FileNotFoundException ex) {
            Logger.getLogger(GenerateKeysAndSecret.class.getName()).log(Level.SEVERE, null, ex);
        } catch (IOException ex) {
            Logger.getLogger(GenerateKeysAndSecret.class.getName()).log(Level.SEVERE, null, ex);
        }
    }

    public void restoreKeysX509(String path, String password) throws FileNotFoundException, IOException, NoSuchAlgorithmException, InvalidKeySpecException, InterruptedException, BadPaddingException {
        try {
            byte[] pubKey = readBytesFromFile(path + File.separatorChar + VotingKeyPair.publicKeyFile);
            // byte[] pubKey = publicKey.getEncoded();

            X509EncodedKeySpec pubKeySpec = new X509EncodedKeySpec(pubKey);
            KeyFactory keyFactory = KeyFactory.getInstance("RSA");
            RSAPublicKey nPublicKey = (RSAPublicKey) keyFactory.generatePublic(pubKeySpec);


            // public key zrobiona 
            CypherModel cyphermodel = new CypherModel();
            System.out.println("haslo  " + password);
            byte[] privKey = cyphermodel.decryptToBytes(readBytesFromFile(path + File.separatorChar + VotingKeyPair.privateKeyFile), password);

            PKCS8EncodedKeySpec privateKeySpec = new PKCS8EncodedKeySpec(
                    privKey);

            RSAPrivateKey nPrivateKey = (RSAPrivateKey) keyFactory.generatePrivate(privateKeySpec);


            privateKey = nPrivateKey;
            publicKey = nPublicKey;

            this.checkCommunication();
        } catch (NoSuchPaddingException ex) {
            Logger.getLogger(GenerateKeysAndSecret.class.getName()).log(Level.SEVERE, null, ex);
        } catch (InvalidKeyException ex) {
            Logger.getLogger(GenerateKeysAndSecret.class.getName()).log(Level.SEVERE, null, ex);
        } catch (IllegalBlockSizeException ex) {
            Logger.getLogger(GenerateKeysAndSecret.class.getName()).log(Level.SEVERE, null, ex);
        }
    }

    byte[] readBytesFromFile(String selectedFile) throws FileNotFoundException, IOException {
        InputStream in = new FileInputStream(selectedFile);
        System.out.println(selectedFile + " rozmiar " + in.available());
        byte[] data = new byte[in.available()];
        in.read(data);
        in.close();
        return data;

    }

    void writeBytesToFile(byte[] data, String fileName) throws FileNotFoundException, IOException {
        FileOutputStream fos = new FileOutputStream(new File(fileName));
        fos.write(data);
        fos.flush();
        fos.close();
    }

    public void checkWritingToFiles() throws NoSuchAlgorithmException, InvalidKeySpecException, IOException { // inne traktujące klucze jako pary bigIntow

        System.out.println("Write public key : " + publicKey.getPublicExponent() + "\n" + publicKey.getModulus());
        System.out.println("Write private key : " + privateKey.getPrivateExponent() + "\n" + privateKey.getModulus());

        storeKeys("keys", kp);

        RSAPublicKey nPublicKey = (RSAPublicKey) this.readPublicKeyFromFile("keys" + File.separatorChar + VotingKeyPair.publicKeyFile);
        RSAPrivateKey nPrivateKey = (RSAPrivateKey) this.readPrivateKeyFrimFile("keys" + File.separatorChar + VotingKeyPair.privateKeyFile);

        System.out.println("Write public key : " + nPublicKey.getPublicExponent() + "\n" + nPublicKey.getModulus());
        System.out.println("Write private key : " + nPrivateKey.getPrivateExponent() + "\n" + nPrivateKey.getModulus());

        System.out.println("New publickey okey : " + nPublicKey.equals(publicKey));
        System.out.println("New privatekey okey : " + nPrivateKey.equals(privateKey));

        System.out.println("Public exponents equal : " + nPublicKey.getPublicExponent().equals(publicKey.getPublicExponent()));
        System.out.println("Private exponents equal : " + nPrivateKey.getPrivateExponent().equals(privateKey.getPrivateExponent()));




    }

    public void storeKeys(String path, KeyPair kp) throws NoSuchAlgorithmException, InvalidKeySpecException, IOException {
        KeyFactory fact = KeyFactory.getInstance("RSA");
        RSAPublicKeySpec pub = fact.getKeySpec(kp.getPublic(),
                RSAPublicKeySpec.class);
        RSAPrivateKeySpec priv = fact.getKeySpec(kp.getPrivate(),
                RSAPrivateKeySpec.class);

        saveToFile(path + File.separatorChar + VotingKeyPair.publicKeyFile, pub.getModulus(),
                pub.getPublicExponent()); // w przyszłości to nie będzie zapisynwane
        saveToFile(path + File.separatorChar + VotingKeyPair.privateKeyFile, priv.getModulus(),
                priv.getPrivateExponent());
    }

    public void saveToFile(String fileName,
            BigInteger mod, BigInteger exp) throws IOException {
        ObjectOutputStream oout = new ObjectOutputStream(
                new BufferedOutputStream(new FileOutputStream(fileName)));
        try {
            oout.writeObject(mod);
            oout.writeObject(exp);
        } catch (Exception e) {
            throw new IOException("Unexpected error", e);
        } finally {
            oout.close();
        }
    }

    PublicKey readPublicKeyFromFile(String keyFileName) throws IOException {

        ObjectInputStream oin = new ObjectInputStream(new BufferedInputStream(new FileInputStream(keyFileName)));
        try {
            BigInteger m = (BigInteger) oin.readObject();
            BigInteger e = (BigInteger) oin.readObject();
            RSAPublicKeySpec keySpec = new RSAPublicKeySpec(m, e);
            KeyFactory fact = KeyFactory.getInstance("RSA");
            PublicKey pubKey = fact.generatePublic(keySpec);
            return pubKey;
        } catch (Exception e) {
            throw new RuntimeException("Spurious serialisation error", e);
        } finally {
            oin.close();
        }
    }

    PrivateKey readPrivateKeyFrimFile(String keyFileName) throws IOException {
        ObjectInputStream oin = new ObjectInputStream(new BufferedInputStream(new FileInputStream(keyFileName)));
        try {
            BigInteger m = (BigInteger) oin.readObject();
            BigInteger e = (BigInteger) oin.readObject();
            RSAPrivateKeySpec keySpec = new RSAPrivateKeySpec(m, e);
            KeyFactory fact = KeyFactory.getInstance("RSA");
            PrivateKey privKey = fact.generatePrivate(keySpec);
            return privKey;
        } catch (Exception e) {
            throw new RuntimeException("Spurious serialisation error", e);
        } finally {
            oin.close();
        }
    }

    public byte[] rsaEncrypt(byte[] data) throws NoSuchAlgorithmException, NoSuchPaddingException, InvalidKeyException, IllegalBlockSizeException, BadPaddingException {
        PublicKey pubKey = publicKey;
        Cipher cipher = Cipher.getInstance("RSA");
        cipher.init(Cipher.ENCRYPT_MODE, pubKey);
        byte[] cipherData = cipher.doFinal(data);
        return cipherData;
    }

    public byte[] rsaDecrypt(byte[] cipherData) throws NoSuchAlgorithmException, NoSuchPaddingException, InvalidKeyException, IllegalBlockSizeException, BadPaddingException {

        Cipher cipher = Cipher.getInstance("RSA");
        cipher.init(Cipher.DECRYPT_MODE, privateKey);
        byte[] data = cipher.doFinal(cipherData);
        return data;
    }

    public void checkCommunication() throws NoSuchAlgorithmException, NoSuchPaddingException, InvalidKeyException, IllegalBlockSizeException, BadPaddingException {
        String secret = "Ala ma kota a kot ma ale ";
        byte[] encrypted = rsaEncrypt(secret.getBytes());

        try {
            System.out.println("Po zaszyfrowaniu : " + new String(encrypted));
        } catch (Exception e) {
            e.printStackTrace();
        }

        byte[] decrypted = rsaDecrypt(encrypted);

        try {
            System.out.println("Po odszyfrowaniu : " + new String(decrypted));
        } catch (Exception e) {
            e.printStackTrace();
        }



    }

    public byte[] getPublicKeyInBytes() {
        if (publicKey == null) {
            return null;
        } else {
            return publicKey.getEncoded();
        }
    }
}
